#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @File   : get_anchors.py
# @Author : yuanwenjin
# @Mail   : xxxx@mail.com
# @Date   : 2020/07/16 09:34:26
# @Docs   : 获取 anchor 大小
'''

import numpy as np
from PIL import Image
import random

def cas_iou(box, cluster):
    x = np.minimum(cluster[:,0],box[0])
    y = np.minimum(cluster[:,1],box[1])

    intersection = x * y
    area1 = box[0] * box[1]

    area2 = cluster[:,0] * cluster[:,1]
    iou = intersection / (area1 + area2 -intersection)

    return iou

def avg_iou(box, cluster):
    return np.mean([np.max(cas_iou(box[i],cluster)) for i in range(box.shape[0])])


def kmeans(box, k):
    # 取出一共有多少框
    row = box.shape[0]
    
    # 每个框各个点的位置
    distance = np.empty((row, k))
    
    # 最后的聚类位置
    last_clu = np.zeros((row,))

    np.random.seed()

    # 随机选5个当聚类中心
    cluster = box[np.random.choice(row, k, replace=False)]
    # cluster = random.sample(row, k)
    while True:
        # 计算每一行距离五个点的iou情况。
        for i in range(row):
            distance[i] = 1 - cas_iou(box[i], cluster)
        
        # 取出最小点
        near = np.argmin(distance, axis=1)

        if (last_clu == near).all():
            break
        
        # 求每一个类的中位点
        for j in range(k):
            cluster[j] = np.median(box[near == j], axis=0)

        last_clu = near

    return cluster

def load_data(path):
    data = []

    for annotation_file in path:
        with open(annotation_file) as f:
            image_labels = f.readlines()
            for img_label in image_labels:
                line = img_label.split()
                image_path = line[0]
                image = Image.open(image_path)
                width, height = image.size
                boxes = line[1:]
                for box in boxes:
                    rect = box.split(',')
                    left = float(rect[0]) / width
                    top = float(rect[1]) / height
                    right = float(rect[2]) / width
                    bottom = float(rect[3]) / height
                    data.append([right-left, bottom-top])

    return np.array(data)


if __name__ == '__main__':

    # 生成 anchors
    SIZE = 608 #　根据训练参数调整
    anchors_num = 9
    # 载入数据集, 格式为: imagepath x1,y1,x2,y2,cls_idx x1,y1,x2,y2,cls_idx ...
    path = ['/home/yuanwenjin/get_samples/papilla/for_yolov4_keras/papilla_label_20200616.txt',
            '/home/yuanwenjin/get_samples/papilla/for_yolov4_keras/papilla_label_20200713.txt']
    
    # 存储格式为转化为比例后的width, height
    data = load_data(path)
    
    # 使用k聚类算法
    out = kmeans(data,anchors_num)
    out = out[np.argsort(out[:,0])]
    print('acc:{:.2f}%'.format(avg_iou(data, out) * 100))
    print(out*SIZE)
    data = out*SIZE
    f = open("papilla_anchors.txt", 'w')
    row = np.shape(data)[0]
    for i in range(row):
        if i == 0:
            x_y = "%d,%d" % (data[i][0], data[i][1])
        else:
            x_y = ", %d,%d" % (data[i][0], data[i][1])
        f.write(x_y)
    f.close()
